package no.ffi.tpm.crypto;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.MGF1ParameterSpec;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.OAEPParameterSpec;
import javax.crypto.spec.PSource;
import javax.security.auth.x500.X500Principal;



/**
 *
 * @author anders
 */
public abstract class CryptoOperations {

    protected PublicKey myPubKey;
    protected PrivateKey myPrivKey;
    public static final String AES="AES";
    public static final String symmetricAlg = "AES/CBC/PKCS5Padding";
    public static final String asymmetricAlgTPM = "RSA/ECB/OAEPWithSHA1AndMGF1Padding";
    public static final String asymmetricAlgPCA="RSA/ECB/PKCS1Padding";
    public static final String signatureAlgorithm = "sha1withrsa";
    
 
  
    public static EncryptedObject encrypt(Serializable object, PublicKey pubKey) throws Exception {
        return encrypt(object,pubKey,asymmetricAlgTPM);
    }
    
    /*
     * Initializes the engine. Parameters must be supplied in some manner
     */
    //public abstract boolean initialize();
    /*
     * RSA encrypts the buffer with the given secret key, then RSA encrypts the
     * secret key. "Serializable" must be replaced with a more tailored class
     */
     public static EncryptedObject encrypt(Serializable object, PublicKey pubKey,String mode) throws Exception {
        // Generate AES symmetric key and encrypt byte array
        SecretKey secretKey = generateSecretKey();//kg.generateKey(); 
        byte[] encryptedBuffer = encrypt(object,secretKey);
        byte[] aesKeyEnc=secretKey.getEncoded();
        // Now encrypt the secret key with rsa
        byte[] buffer=null;
        int preamble=0;
        Cipher cip = Cipher.getInstance(mode);
        //If we have to encrypt for a TPM with a binding key, we need a special preamble
        if (mode.equalsIgnoreCase(asymmetricAlgTPM)) {
            OAEPParameterSpec oaepSpec = new OAEPParameterSpec("SHA1", "MGF1", new MGF1ParameterSpec("SHA1"), new PSource.PSpecified("TCPA".getBytes("ASCII")));
            cip.init(Cipher.ENCRYPT_MODE, pubKey, oaepSpec);
            //must prepare payload for TPM=version||TPM_CONSTANT
            byte[] prolog = new byte[]{(byte) ((short) 1), (byte) ((short) 1), (byte) ((short) 0), (byte) ((short) 0), (byte) ((short) 0x02)}; //TPM version
            byte[] TPMcompatibleBuffer = new byte[prolog.length + aesKeyEnc.length];
            System.arraycopy(prolog, 0, TPMcompatibleBuffer, 0, prolog.length);
            System.arraycopy(aesKeyEnc, 0, TPMcompatibleBuffer, prolog.length, aesKeyEnc.length);
            buffer=TPMcompatibleBuffer;
            preamble=5;
        } else {
            //temporary hack until cryptooperations is refactored properly
            cip.init(Cipher.ENCRYPT_MODE, pubKey);
            buffer = new byte[8 + aesKeyEnc.length];
            System.arraycopy(aesKeyEnc, 0, buffer, 8, aesKeyEnc.length);
            preamble=8;
        }
        
        
        byte[] encryptedKey = cip.doFinal(buffer);
        return new EncryptedObject(mode,preamble,encryptedBuffer, encryptedKey, secretKey);
    }
     
     /*
     * RSA decrypts the buffer. The symmetric key is included in the buffer
     */
    public  abstract byte[] decrypt(EncryptedObject encrypted) throws Exception ;
    
    /*
     * Sign the object with the given private key
     */
    public abstract SignedObject sign(Serializable object) throws Exception;
    /*
     * Verify the signature of the signed object
     */
    public static boolean verify(SignedObject so, PublicKey pubKey) throws Exception {
        Signature s = Signature.getInstance(signatureAlgorithm);
        s.initVerify(pubKey);
        s.update(so.getObjectAsBytes());
        return s.verify(so.getSignature());
    }
    /*
     * Issue a X509 certificate regarding a public key
     */
    abstract X509Certificate certifyPublicKey() throws Exception;
    /*
     * Build a request for an Identity statement
     */
   // public abstract Serializable getISrequest(String password,byte[] nonce) ;
    public abstract Serializable getISrequest() ;
    /*
     * Return the public key of this node
     */
    public abstract PublicKey getPubKey();
    
    /*
     * Generate a symmetric key
     */
    public final static SecretKey generateSecretKey() {
        try {
            KeyGenerator kg = KeyGenerator.getInstance(AES);
            return kg.generateKey();
        } catch (NoSuchAlgorithmException e) { e.printStackTrace(); }
        return null;
    }
    /*
     * Encrypt with a symmetric key
     */
    public final static byte[] encrypt(Serializable object,SecretKey key) {
        try {
            Cipher cip = Cipher.getInstance(symmetricAlg);
            cip.init(Cipher.ENCRYPT_MODE, key);
            byte[] encrBuff=cip.doFinal(serialize(object));
            //in CBC mode an initialization vector is generated for encryption
            //and it must be sent in clear to the receiver to perform decryption
            //Here we concatenate it to the cipher as the first 16 bytes of the output
            //Iv is automatically generated by the Cipher object
            byte[] iv=cip.getIV();
            byte[] out=new byte[iv.length+encrBuff.length];
            System.arraycopy(iv, 0, out, 0, iv.length);
            System.arraycopy(encrBuff, 0, out, iv.length, encrBuff.length);
            return out;
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) { e.printStackTrace(); }
        return null;
    }
    /*
     * Decrypt with symmetric key
     */
    public final static byte[] decrypt(byte[] cipher, SecretKey key) {
        try {
            //SecretKey sessKey = new SecretKeySpec(key.getEncoded(),symmetricAlg);
            Cipher cip = Cipher.getInstance(symmetricAlg);
            byte[] iv=new byte[16];
            System.arraycopy(cipher, 0, iv, 0, 16);//extract the IV
            byte[] encBuffer=new byte[cipher.length-16];
            System.arraycopy(cipher, 16, encBuffer,0,cipher.length-16);//extract the actual encrypted buffer
            cip.init(Cipher.DECRYPT_MODE,key, new IvParameterSpec(iv));
            byte[] plaintext = cip.doFinal(encBuffer);
            return plaintext;
           // return deserialize(plaintext); Not possible when decrypting TPM encrypted data
        } catch (Exception e) { e.printStackTrace(); }
        return null;
    }
    public static byte[] serialize(Serializable o) {
        try {
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            ObjectOutputStream oos = new ObjectOutputStream(bos);
            oos.writeObject(o);
            byte[] b = bos.toByteArray();
            return b;
        } catch (IOException e) { e.printStackTrace(); }
        return null;
    }
    
    public static Serializable deserialize(byte[] buf) {
        try {
            ByteArrayInputStream bis = new ByteArrayInputStream(buf);
            ObjectInputStream ois = new ObjectInputStream(bis);
            return (Serializable) ois.readObject();
        } catch (ClassNotFoundException|IOException e) { e.printStackTrace(); }
        return null;
    }
    
    static char[] hexVal = {'0','1','2','3','4','5','6','7','8',
                        '9','a','b','c','d','e','f'};

    public static String toHexString(byte[] b) {
        StringBuffer sb = new StringBuffer();
        for (int i=0;i<b.length;i++) {
            //System.out.print(b[i]>>4);
            sb.append(hexVal[b[i]>>4&0xf]);
            sb.append(hexVal[b[i]&0xf]);
        }
        return new String(sb);
    }

    //The SHA-1 hash of a public key as calculated by the TPM
    public static byte[] pubKeytoSHA1(PublicKey key) {
        byte[] calculatedDigest = null;

        //Extract the modulus of the key 

        byte[] modTemp = ((RSAPublicKey) key).getModulus().toByteArray();
        //Remove the first byte
        byte[] mod = new byte[modTemp.length - 1];
        System.arraycopy(modTemp, 1, mod, 0, modTemp.length - 1);
        //Create a digest of it 

        calculatedDigest = SHA1(mod);


        return calculatedDigest;
    }
     
    public static byte[] SHA1(byte[] buffer) {
        byte[] calculatedDigest=null;
        try {
            MessageDigest sha = MessageDigest.getInstance("SHA-1");
            sha.update(buffer);
            calculatedDigest = sha.digest();
        } catch (NoSuchAlgorithmException ex) {
            Logger.getLogger(CryptoOperations.class.getName()).log(Level.SEVERE, null, ex);
        }
        return calculatedDigest;
    }

    /**
     * @return the myDistinguishedName
     */
    public abstract String getMyDistinguishedName();    
  
    public abstract X509Certificate getTrustAnchorCert();

    
    public abstract PublicKey getTrustAnchor() ;
    
    
    

    
    public abstract X500Principal getSubjectPrincipal() ;
    
}
